package com.zhangdy.agent.common.controller;


import com.zhangdy.agent.common.dto.MethodPoint;
import com.zhangdy.agent.common.dto.LoadAgent;
import com.zhangdy.agent.common.dto.LoadClass;
import com.zhangdy.agent.common.dto.ReturnT;
import com.zhangdy.agent.common.utils.DynamicCompilerUtil;
import com.zhangdy.agent.common.utils.GsonTool;
import com.zhangdy.agent.common.utils.HotSwapClient;
import javassist.ClassPool;
import javassist.CtClass;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Method;

@RestController
public class AgentController {

    private static Logger logger = LoggerFactory.getLogger(AgentController.class);


    @PostMapping(value = "loadAgent")
    public ReturnT loadAgent(@RequestBody LoadAgent loadAgent){
        Object virtualMachine = null;
        Class<?> virtualMachineClazz = null;
        try {

            logger.info("loadAgent:{}",  GsonTool.toJson(loadAgent));
            String name = ManagementFactory.getRuntimeMXBean().getName();
            // get pid
            String pid = name.split("@")[0];
            String clazz = "com.sun.tools.attach.VirtualMachine";
            virtualMachineClazz = Class.forName(clazz);

            Method attachMethod = virtualMachineClazz.getMethod("attach", String.class);
            Method loadAgentMethod = virtualMachineClazz.getMethod("loadAgent", String.class);

            virtualMachine = attachMethod.invoke(null, pid);
            loadAgentMethod.invoke(virtualMachine, loadAgent.getPath());
        }catch (Exception e) {
            logger.error("loadAgent:{}",  GsonTool.toJson(loadAgent), e);

            return ReturnT.FAIL;
        } finally {
            try {
                Method loadAgentMethod = virtualMachineClazz.getMethod("detach");
                loadAgentMethod.invoke(virtualMachine);
            }catch (Exception e) {

            }

        }


        return ReturnT.SUCCESS;
    }


    @PostMapping(value = "loadClass")
    public ReturnT loadClass(@RequestBody LoadClass loadClass){
        try {
            logger.info("loadClass:{}",  GsonTool.toJson(loadClass));

            String className = loadClass.getClassName();
            String sourceCode = loadClass.getSourceCode();
            DynamicCompilerUtil.arthasDynamicCompiler(className, sourceCode);
            String path = DynamicCompilerUtil.CUSTOMIZE_CLZ_DIR.concat("/").concat(StringUtils.replace(className, ".", File.separator) + ".class");
            String s = DynamicCompilerUtil.loadByteCode(className,path);
            return new ReturnT(s);
        }catch (Exception e) {
            logger.error("loadClass:{}",  GsonTool.toJson(loadClass), e);
            return ReturnT.FAIL;
        }
    }


    @PostMapping(value = "redefineClass")
    public ReturnT redefineClass(@RequestBody MethodPoint point){
        try {
            logger.info("redefineClass:{}",  GsonTool.toJson(point));

            ClassPool pool = ClassPool.getDefault();
            String targetClassName = point.getPackageName() + "." + point.getClassName();
            CtClass ctClass = DynamicCompilerUtil.replaceByteCode(point, pool);
            HotSwapClient.redefine(Class.forName(targetClassName), ctClass);
            ctClass.defrost();

        }catch (Exception e) {
            logger.error("redefineClass:{}",  GsonTool.toJson(point), e);
            return ReturnT.FAIL;

        }
        return ReturnT.SUCCESS;
    }

}
